import wave
import math
import os
import numpy

Factor = 1024
FactorPro = 1150

def Calculate_Cn(input):
    if (input > 7 * Factor):
        coding = 15
    elif (input < -8 * Factor):
        coding = 0
    else:
        coding = math.floor(input / Factor) + 9
    return coding

def Calculate_Cn_Pro(input):
    if (input > 7 * FactorPro):
        coding = 15
    elif (input < -8 * FactorPro):
        coding = 0
    else:
        coding = math.floor(input / FactorPro) + 9
    return coding

def DPCM(waveData, frames):
    diff = numpy.zeros((frames, 1))
    Xbar = numpy.zeros((frames, 1))
    Cn = numpy.zeros((frames, 1), dtype=int)
    Cn[0] = diff[0] = Xbar[0] = waveData[0]
    for i in range(1, frames):
        diff[i] = waveData[i] - Xbar[i-1]
        Cn[i] = Calculate_Cn(diff[i])
        Xbar[i] = Xbar[i-1] + (Cn[i] - 8) * Factor
    print ('DPCM SNR:')
    Calculate(waveData, Xbar)
    print (' ')
    return Cn

def DPCM_Pro(waveData, frames):
    diff = numpy.zeros((frames, 1))
    Xbar = numpy.zeros((frames, 1))
    Cn = numpy.zeros((frames, 1), dtype=int)
    Cn[0] = diff[0] = Xbar[0] = waveData[0]
    for i in range(0, len(waveData)):
        if (i % 2 == 1) and  (i != frames):
            waveData[i] = (waveData[i-1] + waveData[i+1]) / 2
        else:
            continue
    for i in range(1, frames):
        diff[i] = waveData[i] - Xbar[i-1]
        Cn[i] = Calculate_Cn(diff[i])
        Xbar[i] = Xbar[i-1] + (Cn[i] - 8) * FactorPro
    print ('DPCM Pro Version SNR:')
    Calculate(waveData, Xbar)
    return Cn

def Decode(dpcData, frames):
    pcmData = []
    pcmData += [dpcData[0]]
    for i in range(1, frames):
        pcmData += [(pcmData[i-1] + (dpcData[i] - 8) * Factor)]
    return pcmData

def Decode_Pro(dpcData, frames):
    pcmData = []
    pcmData += [dpcData[0]]
    for i in range(1, frames):
        pcmData += [(math.log(abs(pcmData[i-1])+0.1) + (dpcData[i] - 8) * Factor - 0.1)]
    return pcmData

def SaveDpcFile(Cn):
    path = os.path.abspath('.')
    dpcPath = path + '/dpc/'
    dpcFile = dpcPath + '1.dpc'
    fileWriter = open(dpcFile, 'wb')
    fileWriter.write(bytes((0xFFFF & int(Cn[0])).to_bytes(2, byteorder="big")))
    for i in range(1, len(Cn)):
        if (i % 2 == 1) and (i != len(Cn) - 1):
            tmp = 0xFF & int(Cn[i])
            tmp = tmp << 4
            tmp = tmp | int(Cn[i+1])
            fileWriter.write(bytes((0xFF & tmp).to_bytes(1, byteorder="big")))
        elif (i % 2 ==0) and (i != len(Cn) - 1):
            continue
        else:
            tmp = 0xF & int(Cn[i])
            tmp = tmp << 4
            fileWriter.write(bytes((0xFF & tmp).to_bytes(1, byteorder="big")))
    fileWriter.close()

def Calculate(waveData, DPCM_Data):
    x = 0
    y = 0
    for i in range(0, len(waveData)):
        x = x + DPCM_Data[i] * DPCM_Data[i]
        y = y + (DPCM_Data[i] - waveData[i]) * (DPCM_Data[i] - waveData[i])
    z = 10 * math.log((x / (y * 1.0)), 10)
    print(z)

if __name__ == '__main__':

    # Read file
    path = os.path.abspath('.')
    wavePath = path + '/WAVE/'
    dpcPath = path + '/dpc/1.dpc'
    waveTestFilePath = wavePath + '1.wav'
    waveTestFileReader = wave.open(waveTestFilePath, 'rb')
    nframes = waveTestFileReader.getnframes()
    width = waveTestFileReader.getsampwidth()
    channels = waveTestFileReader.getnchannels()
    rate = waveTestFileReader.getframerate()
    tempWaveData = waveTestFileReader.readframes(nframes)
    waveTestFileReader.close()
    waveData = numpy.fromstring(tempWaveData, dtype=numpy.int16)

    # DPCM
    Cn = DPCM(waveData, nframes)
    # DPCM Pro
    CnPro = DPCM_Pro(waveData, nframes)
    pcmData = Decode(Cn, nframes)
    pcmDataPro = Decode(CnPro, nframes)

    # Save data in .dpc file
    SaveDpcFile(Cn)
    # Save data in .pcm file
    pcmPath = path + '/pcm/'
    pcmFile = pcmPath + '1.pcm'
    pcmProFile = pcmPath + 'pro.pcm'
    file_write = wave.open(pcmFile, 'wb')
    file_write.setframerate(rate*4)
    file_write.setnchannels(channels)
    file_write.setsampwidth(width)
    file_write.writeframes(numpy.array(pcmData).tostring())
    file_write.close()
    file_write = wave.open(pcmProFile, 'wb')
    file_write.setframerate(rate*4)
    file_write.setnchannels(channels)
    file_write.setsampwidth(width)
    file_write.writeframes(numpy.array(pcmDataPro).tostring())
    file_write.close()
